package com.bestvike.stone.spring.context;

import com.bestvike.stone.core.threading.InheritThreadLocal;
import com.bestvike.stone.spring.trace.ISpan;
import com.bestvike.stone.spring.trace.TraceHeaders;
import com.bestvike.stone.spring.trace.TracerUtils;
import org.slf4j.MDC;

import java.util.ArrayDeque;
import java.util.Deque;

/**
 * Created by 许崇雷 on 2018-01-26.
 */
public final class TraceContext {
    //线程本地存储
    private static final ThreadLocal<ISpan> current = InheritThreadLocal.withInitial(() -> ISpan.EMPTY);
    private static final ThreadLocal<Deque<ISpan>> contexts = ThreadLocal.withInitial(() -> {
        Deque<ISpan> stack = new ArrayDeque<>();
        stack.push(current.get());
        return stack;
    });

    //region property

    public static boolean exists() {
        return current.get() != ISpan.EMPTY;
    }

    private static ISpan getSpan() {
        return current.get();
    }

    public static String getTraceId() {
        return getSpan().getTraceId();
    }

    public static String getSpanId() {
        return getSpan().getSpanId();
    }

    public static String getParentSpanId() {
        return getSpan().getParentSpanId();
    }

    public static String getTraceSpanId() {
        return getSpan().getTraceSpanId();
    }

    //endregion

    public static void init() {
        ISpan span = TracerUtils.createSpan();
        Deque<ISpan> stack = contexts.get();
        stack.clear();
        stack.push(span);
        current.set(span);
        MDC.put(TraceHeaders.TRACE_ID_NAME, span.getTraceId());
        MDC.put(TraceHeaders.SPAN_ID_NAME, span.getSpanId());
        MDC.put(TraceHeaders.PARENT_SPAN_ID_NAME, span.getParentSpanId());
        MDC.put(TraceHeaders.TRACE_SPAN_ID_NAME, span.getTraceSpanId());
    }

    public static void reset() {
        contexts.get().clear();
        current.set(ISpan.EMPTY);
        MDC.remove(TraceHeaders.TRACE_ID_NAME);
        MDC.remove(TraceHeaders.SPAN_ID_NAME);
        MDC.remove(TraceHeaders.PARENT_SPAN_ID_NAME);
        MDC.remove(TraceHeaders.TRACE_SPAN_ID_NAME);
    }

    public static void beginSpan() {
        ISpan span = TracerUtils.createContinueSpan(getSpan());
        contexts.get().push(span);//must call 'contexts.get()' in sub thread to init stack when first beginSpan().
        current.set(span);
        MDC.put(TraceHeaders.TRACE_ID_NAME, span.getTraceId());
        MDC.put(TraceHeaders.SPAN_ID_NAME, span.getSpanId());
        MDC.put(TraceHeaders.PARENT_SPAN_ID_NAME, span.getParentSpanId());
        MDC.put(TraceHeaders.TRACE_SPAN_ID_NAME, span.getTraceSpanId());
    }

    public static void endSpan() {
        Deque<ISpan> stack = contexts.get();
        if (stack.size() < 2)
            return;
        stack.pop();
        ISpan span = stack.getFirst();
        current.set(span);
        MDC.put(TraceHeaders.TRACE_ID_NAME, span.getTraceId());
        MDC.put(TraceHeaders.SPAN_ID_NAME, span.getSpanId());
        MDC.put(TraceHeaders.PARENT_SPAN_ID_NAME, span.getParentSpanId());
        MDC.put(TraceHeaders.TRACE_SPAN_ID_NAME, span.getTraceSpanId());
    }
}
